/**
 * @file
 * @author dousha 
 * @date 2020/11/28.
 */

#include <string.h>
#include "serial_state_machine.h"
#include "../Peripherals/serial.h"
#include "../../smx/sm3.h"
#include "../Peripherals/entropy.h"
#include "convert.h"
#include "constants.h"
#include "menu_state_machine.h"
#include "../Peripherals/eeprom.h"
#include "crypto_core.h"
#include "saved_secrets.h"
#include "hmac.h"

static uint8_t state = WAIT_UNLOCK, function = 0, param = 0, more = 0;
static uint64_t length = 0, offset = 0, remainingLength = 0, processedLength = 0, targetLength = 0, bitLength = 0;
static uint8_t txBuffer[512] = {0};
static uint8_t stateFlag = 0, secretEntry = 0;

void reply_bad_command() {
	serial_write_string("?\r\n");
}

void reply_success() {
	serial_write_string("!\r\n");
}

void reply_entity_too_large() {
	serial_write_string("#\r\n");
}

void swap_endian_32(uint8_t *data, size_t len) {
	uint8_t t;
	for (size_t i = 0; i < len; i += 4) {
		t = data[i];
		data[i] = data[i + 3];
		data[i + 3] = t;
		t = data[i + 1];
		data[i + 1] = data[i + 2];
		data[i + 2] = t;
	}
}

void evaluate_data(const uint8_t *data) {
	uint32_t *sm3_out;
	switch (function) {
		case FUNCTION_SM2:
			break;
		case FUNCTION_SM3:
			if (stateFlag == 0) {
				stateFlag = 1;
				reply_success();
				sm3_init();
				bitLength = length * 8;
				targetLength = sm3_get_padded_byte_length(bitLength);
				processedLength = 0;
			}
			// sm3_update always takes 64 bytes
			sm3_update(data, processedLength, bitLength);
			processedLength += 64;
			if (remainingLength == 0) {
				while (processedLength < targetLength) {
					sm3_update(data, processedLength, bitLength);
					processedLength += 64;
				}
				sm3_out = sm3_finalize();
				memcpy(txBuffer, sm3_out, 32);
				swap_endian_32(txBuffer, 32);
				serial_write_hex(txBuffer, 32);
			}
			break;
		case FUNCTION_SM4_ENCRYPT:
			if (param & 0x01u) {
				// user key
				// TODO
			} else {
				// user data block
				if (stateFlag == 0) {
					reply_success();
					stateFlag = 1;
					prepare_encrypt();
					targetLength = (length + 15) & 0xfffffffffffffff0u;
				}
				if (remainingLength > 0) {
					for (uint16_t i = 0; i < SERIAL_RX_SIZE; i += 16) {
						for (uint8_t j = 0; j < 16; j++) {
							txBuffer[j] = i + j < length ? data[i + j] : 0;
						}
						encrypt_block(txBuffer, txBuffer + 16);
						serial_write_hex(txBuffer + 16, 16);
					}
				} else {
					for (uint8_t i = 0; i < 16; i++) {
						txBuffer[i] = i < length ? data[i] : 0;
					}
					encrypt_block(txBuffer, txBuffer + 16);
					serial_write_hex(txBuffer + 16, 16);
				}
			}
			break;
		case FUNCTION_SM4_DECRYPT:
			if (param & 0x01u) {
				// user key
				// TODO
			} else {
				if (stateFlag == 0) {
					reply_success();
					stateFlag = 1;
					prepare_decrypt();
					targetLength = (length + 15) & 0xfffffffffffffff0u;
				}
				if (remainingLength > 0) {
					for (uint16_t i = 0; i < SERIAL_RX_SIZE; i += 16) {
						for (uint8_t j = 0; j < 16; j++) {
							txBuffer[j] = i + j < length ? data[i + j] : 0;
						}
						decrypt_block(txBuffer, txBuffer + 16);
						serial_write_bytes(txBuffer + 16, 16);
						processedLength += 16;
					}
				} else {
					for (uint8_t i = 0; i < 16; i++) {
						txBuffer[i] = i < length ? data[i] : 0;
					}
					decrypt_block(txBuffer, txBuffer + 16);
					serial_write_bytes(txBuffer + 16, 16);
					processedLength += 16;
				}
			}
			break;
		case FUNCTION_SECRET_READ:
			if (param & 0x80u) {
				// seek
				offset = secret_seek((char *) data, length);
				if (offset == 0xff) {
					reply_bad_command();
					return;
				}
			} else {
				// offset
				offset = data[0];
				secret_index_load((uint8_t) offset);
				if (strlen((const char *) secret_current_get()->name) == 0) {
					reply_bad_command();
					return;
				}
			}
			reply_success();
			secret_get_prepare(offset);
			do {
				more = secret_read_page(txBuffer);
				serial_write_bytes(txBuffer, 32);
			} while (more);
			break;
		case FUNCTION_SECRET_STORE:
			switch (stateFlag) {
				case 0: // name
					// param: otxx_xxxx
					if (param & 0x80u) {
						// overwrite entry
						secretEntry = param & 0x3fu;
						secret_set_prepare();
					} else {
						// add entry
						secretEntry = secret_allocate();
						if (secretEntry != 0xff) {
							secret_set_prepare();
						} else {
							reply_bad_command();
							stateFlag = 0xff;
							return;
						}
					}
					if (param & 0x40u) {
						secret_current_get()->metadata |= SECRET_META_HEX;
					}
					memcpy((char *) secret_current_get()->name, data, length > 16 ? 15 : length);
					secret_current_get()->name[15] = 0;
					stateFlag = 1;
					break;
				case 1: // body
					if (remainingLength == 0) {
						for (uint32_t i = 0; i < 32; i++) {
							txBuffer[i] = processedLength + i < length ? data[i] : 0;
						}
						secret_write_page(txBuffer);
						stateFlag = 2;
					} else {
						for (uint32_t i = 0; i < SERIAL_RX_SIZE; i += 32) {
							memcpy(txBuffer, data, 32);
							secret_write_page(txBuffer);
							processedLength += 32;
						}
					}
					break;
				case 2: // time
					secret_current_get()->length = length;
					secret_set_time(get_u16(data), get_u16(data + 2), get_u16(data + 4), get_u16(data + 6),
									get_u16(data + 8));
					secret_index_save(secretEntry);
					reply_success();
					stateFlag = 0;
					break;
				default:
					break;
			}
			break;
		case FUNCTION_HMAC:
			if (stateFlag == 0) {
				if (param & 0x40u) {
					offset = 0;
					do {
						more = secret_read_page(txBuffer + offset);
						offset += 16;
					} while (more);
					offset = 0;
					hmac_prepare(txBuffer, secret_current_get()->length, length);
					goto hmac_msg_processing;
				} else if (param & 0x80u) {
					hmac_prepare(secret_master_secret(), 16, length);
					goto hmac_msg_processing;
				} else {
					if (length > 512) {
						return;
					} else {
						processedLength = 0;
						if (remainingLength == 0) {
							memcpy(txBuffer, data, length);
							stateFlag = 2;
						} else {
							memcpy(txBuffer, data, SERIAL_RX_SIZE);
							stateFlag = 1;
							processedLength += SERIAL_RX_SIZE;
						}
					}
				}
			} else if (stateFlag == 1) {
				// key, continuation
				if (remainingLength == 0) {
					memcpy(txBuffer + processedLength, data, length - processedLength);
					processedLength = length;
					stateFlag = 2;
				} else {
					memcpy(txBuffer, data, SERIAL_RX_SIZE);
					stateFlag = 1;
					processedLength += SERIAL_RX_SIZE;
				}
			} else if (stateFlag == 2) {
				// message
				hmac_prepare(txBuffer, processedLength, length);
				hmac_msg_processing:
				processedLength = 0;
				if (remainingLength == 0) {
					memset(txBuffer, 0, 64);
					memcpy(txBuffer, data, length);
					hmac_update(txBuffer);
					memcpy(txBuffer + 64, hmac_finalize(), 32);
					serial_write_hex(txBuffer + 64, 32);
				} else {
					hmac_update(data);
					processedLength += 64;
					stateFlag = 3;
				}
			} else {
				// message
				if (remainingLength == 0) {
					memset(txBuffer, 0, 64);
					memcpy(txBuffer, data, length - processedLength);
					hmac_update(txBuffer);
					memcpy(txBuffer + 64, hmac_finalize(), 32);
					serial_write_hex(txBuffer + 64, 32);
					stateFlag = 0;
				} else {
					hmac_update(data);
					processedLength += 64;
				}
			}
			break;
		case FUNCTION_EEPROM_INSPECT:
			offset = get_u16(data);
			length = get_u16(data + 2);
			eeprom_read(offset, txBuffer, length);
			serial_write_hex(txBuffer, length);
			break;
		default:
			break;
	}
}

size_t next_part_length() {
	switch (function) {
		case FUNCTION_SM4_ENCRYPT:
		case FUNCTION_SM4_DECRYPT:
			state = WAIT_DATA;
			return 8; // data fragments
		case FUNCTION_SECRET_STORE:
			switch (stateFlag) {
				case 1: // name
					state = WAIT_MULTIPART_LENGTH;
					return 8;
				case 2: // body
					state = WAIT_DATA;
					return 10;
				default:
					state = WAIT_PREAMBLE;
					return 1;
			}
		case FUNCTION_HMAC:
			if (length > 512) {
				reply_entity_too_large();
				state = WAIT_PREAMBLE;
				return 1;
			} else {
				// provided key received, message now
				state = WAIT_LENGTH;
				return 8;
			}
		default:
			state = WAIT_PREAMBLE;
			return 1;
	}
}

size_t evaluate_param() {
	switch (function) {
		case FUNCTION_SM4_ENCRYPT:
		case FUNCTION_SM4_DECRYPT:
			// mmmm ___x
			if ((param & 1u) == 0) {
				// use master secret
				state = WAIT_LENGTH;
			} else {
				// use provided key
				state = WAIT_MULTIPART_LENGTH;
			}
			return 8;
		case FUNCTION_SECRET_READ:
			if (param & 0x80u) {
				// seek
				if ((param & 0x0fu) == 0) {
					reply_bad_command();
					state = WAIT_PREAMBLE;
					return 1;
				} else {
					state = WAIT_DATA;
					return (param & 0x0fu);
				}
			} else {
				// offset
				state = WAIT_DATA;
				return 1;
			}
		case FUNCTION_SECRET_STORE:
			state = WAIT_MULTIPART_LENGTH;
			return 8;
		case FUNCTION_HMAC:
			if (param & 0x80u) {
				// use master secret as hmac key
				state = WAIT_LENGTH;
			} else if (param & 0x40u) {
				// use a saved secret
				state = WAIT_LENGTH;
				if (secret_index_is_valid(param & 0x3fu)) {
					secret_get_prepare(param & 0x3fu);
					if (secret_current_get()->length > 512) {
						reply_entity_too_large();
						state = WAIT_PREAMBLE;
						return 1;
					}
				} else {
					reply_bad_command();
					state = WAIT_PREAMBLE;
					return 1;
				}
			} else {
				// use provided key
				state = WAIT_MULTIPART_LENGTH;
			}
			return 8;
		default:
			reply_bad_command();
			state = WAIT_PREAMBLE;
			return 1;
	}
}

size_t evaluate_function(uint8_t fn) {
	switch (fn) {
		case 'V':
			serial_write_string(FIRMWARE_VERSION);
			state = WAIT_PREAMBLE;
			return 1;
		case 'H':
			state = WAIT_LENGTH;
			function = FUNCTION_SM3;
			return 8;
		case 'E':
			state = WAIT_PARAMETER;
			function = FUNCTION_SM4_ENCRYPT;
			return 1;
		case 'D':
			state = WAIT_PARAMETER;
			function = FUNCTION_SM4_DECRYPT;
			return 1;
		case 'A':
			state = WAIT_PARAMETER;
			function = FUNCTION_HMAC;
			return 1;
		case 'R':
			state = WAIT_PREAMBLE;
			serial_write_u8(entropy_read());
			return 1;
		case 'I':
			state = WAIT_DATA;
			function = FUNCTION_EEPROM_INSPECT;
			return 4;
		case '<':
			state = WAIT_PARAMETER;
			function = FUNCTION_SECRET_STORE;
			return 1;
		case '>':
			state = WAIT_PARAMETER;
			function = FUNCTION_SECRET_READ;
			return 1;
		default:
			reply_bad_command();
			state = WAIT_PREAMBLE;
			return 1;
	}
}

size_t serial_state_machine_handle_input(const uint8_t *bytes) {
	switch (state) {
		case WAIT_UNLOCK:
			if (!is_device_unlocked()) {
				reply_bad_command();
				return 1;
			}
			state = WAIT_PREAMBLE;
			// falling through intended
		case WAIT_PREAMBLE:
			if (bytes[0] == '$') {
				state = WAIT_CS;
				return 1;
			} else {
				reply_bad_command();
				return 1;
			}
		case WAIT_CS:
			return evaluate_function(bytes[0]);
		case WAIT_PARAMETER:
			param = bytes[0];
			return evaluate_param();
		case WAIT_LENGTH:
		case WAIT_MULTIPART_LENGTH:
			if (state == WAIT_LENGTH) {
				state = WAIT_DATA;
			} else {
				state = WAIT_MULTIPART_DATA;
			}
			length = get_u64(bytes);
			if (length == 0) {
				reply_bad_command();
				state = WAIT_PREAMBLE;
				return 1;
			} else {
				remainingLength = length;
				return length > SERIAL_RX_SIZE ? SERIAL_RX_SIZE : length;
			}
		case WAIT_DATA:
		case WAIT_MULTIPART_DATA:
			if (remainingLength > SERIAL_RX_SIZE) {
				remainingLength -= SERIAL_RX_SIZE;
				evaluate_data(bytes);
				return remainingLength > SERIAL_RX_SIZE ? SERIAL_RX_SIZE : remainingLength;
			} else {
				remainingLength = 0;
				evaluate_data(bytes);
				if (state == WAIT_DATA) {
					stateFlag = 0;
					state = WAIT_PREAMBLE;
					return 1;
				} else {
					return next_part_length();
				}
			}
		default:
			state = WAIT_PREAMBLE;
			return 1;
	}
}
